// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/sparksql/specialforms/get_struct_field.h>
#include <pollux/expression/constant_expr.h>
#include <pollux/vector/complex_vector.h>

namespace kumo::pollux::functions::sparksql {
    namespace {
        // Returns the value of nested subfield in the input struct.
        // The input must be of row type and nested complex type is allowed.
        // The subfield position is specified by 'ordinal'.
        // If 'ordinal' is negative or greater than the children size of input,
        // exception is thrown.
        class GetStructFieldFunction : public exec::VectorFunction {
        public:
            explicit GetStructFieldFunction(int32_t ordinal) : ordinal_(ordinal) {
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /*resultType*/,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                exec::LocalDecodedVector decoded(context, *args[0], rows);
                auto rowData = decoded->base()->as<RowVector>();
                POLLUX_USER_CHECK_LT(
                    ordinal_,
                    rowData->childrenSize(),
                    "Invalid ordinal. Should be smaller than the children size of input row vector.");
                if (decoded->isIdentityMapping()) {
                    result = rowData->childAt(ordinal_);
                } else {
                    result =
                            decoded->wrap(rowData->childAt(ordinal_), *args[0], decoded->size());
                }
            }

        private:
            // The position to select subfield from the struct.
            const int32_t ordinal_;
        };
    } // namespace

    TypePtr GetStructFieldCallToSpecialForm::resolveType(
        const std::vector<TypePtr> & /*argTypes*/) {
        POLLUX_FAIL("GetStructField function does not support type resolution.");
    }

    exec::ExprPtr GetStructFieldCallToSpecialForm::constructSpecialForm(
        const TypePtr &type,
        std::vector<exec::ExprPtr> &&args,
        bool trackCpuUsage,
        const core::QueryConfig & /*config*/) {
        POLLUX_USER_CHECK_EQ(
            args.size(), 2, "get_struct_field expects two arguments.");

        POLLUX_USER_CHECK_EQ(
            args[0]->type()->kind(),
            TypeKind::ROW,
            "The first argument of get_struct_field should be of row type.");

        POLLUX_USER_CHECK_EQ(
            args[1]->type()->kind(),
            TypeKind::INTEGER,
            "The second argument of get_struct_field should be of integer type.");

        auto constantExpr = std::dynamic_pointer_cast<exec::ConstantExpr>(args[1]);
        POLLUX_USER_CHECK_NOT_NULL(
            constantExpr,
            "The second argument of get_struct_field should be constant expression.");
        POLLUX_USER_CHECK(
            constantExpr->value()->is_constant_encoding(),
            "The second argument of get_struct_field should be wrapped in constant vector.");
        auto constant_vector =
                constantExpr->value()->as_unchecked<ConstantVector<int32_t> >();
        POLLUX_USER_CHECK(
            !constant_vector->is_null_at(0),
            "The second argument of get_struct_field should be non-nullable.");

        auto ordinal = constant_vector->value_at(0);

        POLLUX_USER_CHECK_GE(ordinal, 0, "Invalid ordinal. Should be greater than 0.");
        auto getStructFieldFunction =
                std::make_shared<GetStructFieldFunction>(ordinal);

        return std::make_shared<exec::Expr>(
            type,
            std::move(args),
            std::move(getStructFieldFunction),
            exec::VectorFunctionMetadata{},
            kGetStructField,
            trackCpuUsage);
    }
} // namespace kumo::pollux::functions::sparksql
